// Package main 实现了泛型、线程安全的Count-Min Sketch概率数据结构
// Count-Min Sketch（CM Sketch）是一种用于海量流数据频率统计的轻量级结构，
// 通过牺牲“精确性”换取“亚线性内存”和“高效操作”，适用于网络IP统计、日志分析、热门元素TopK等场景
// 核心优势：内存开销O(d×w)、插入/查询时间O(d)、线程安全、支持分布式合并
package main

import (
	"encoding/binary" // 用于种子的字节序列化
	"fmt"             // 用于元素字符串化和错误提示
	"hash/fnv"        // 提供FNV-1a哈希算法（轻量、高效，适合非加密场景）
	"math"            // 用于数学常量（如MaxUint64）
	"math/rand"       // 用于生成随机哈希种子
	"sort"            // 用于TopK功能的排序
	"sync/atomic"     // 提供原子操作，保证线程安全
	"time"            // 用于初始化随机种子（时间戳确保种子唯一性）
)

// CountMinSketch 是Count-Min Sketch概率数据结构的泛型实现
// 支持任意可比较类型（T comparable）的频率统计，通过原子计数器保证多线程安全
type CountMinSketch[T comparable] struct {
	width uint64 // 矩阵列数（对应CM Sketch的w）：每个哈希函数的映射范围为[0, width-1]
	// 列数越多，不同元素哈希碰撞到同一列的概率越低，频率估算误差越小
	depth uint64 // 矩阵行数（对应CM Sketch的d）：等于独立哈希函数的个数
	// 行数越多，“多个哈希函数同时碰撞”的概率越低，误差概率越小
	seeds []uint32 // 哈希函数的独立种子：每个行对应一个种子，确保哈希函数独立性
	// 避免不同行使用相同哈希逻辑导致的系统性碰撞
	counters [][]atomic.Uint64 // 原子计数器矩阵：depth行×width列，存储元素的累计频率
	// 使用atomic.Uint64确保多线程并发插入/查询时计数器操作的原子性，无需显式加锁
}

// NewCountMinSketch 创建默认的Count-Min Sketch实例（使用随机种子）
// 参数：
//
//	width: 矩阵列数（w），必须为正整数
//	depth: 矩阵行数（d），必须为正整数
//
// 返回：
//
//	*CountMinSketch[T]: 初始化后的CM Sketch实例
//
// 注意：
//   - 若width或depth为0，会触发panic（无效参数）
//   - 种子通过当前时间戳初始化，确保每次创建实例的种子不同（除非同一纳秒调用）
func NewCountMinSketch[T comparable](width, depth uint64) *CountMinSketch[T] {
	// 参数校验：列数和行数必须为正（否则矩阵无意义，无法存储计数器）
	if width == 0 || depth == 0 {
		panic("width and depth must be positive")
	}

	// 初始化随机种子：用当前时间戳（纳秒级）确保每次运行生成的种子不同
	rand.Seed(time.Now().UnixNano())
	// 生成depth个独立随机种子（每个种子32位，保证哈希函数独立性）
	seeds := make([]uint32, depth)
	for i := range seeds {
		seeds[i] = rand.Uint32()
	}

	// 初始化原子计数器矩阵：depth行，每行width个计数器
	// atomic.Uint64默认初始值为0，无需额外初始化
	counters := make([][]atomic.Uint64, depth)
	for i := range counters {
		counters[i] = make([]atomic.Uint64, width)
	}

	// 组装并返回CM Sketch实例
	return &CountMinSketch[T]{
		width:    width,
		depth:    depth,
		seeds:    seeds,
		counters: counters,
	}
}

// NewCountMinSketchWithSeeds 创建指定种子的Count-Min Sketch实例（主要用于测试和分布式场景）
// 参数：
//
//	width: 矩阵列数（w），必须为正整数
//	depth: 矩阵行数（d），必须为正整数
//	seeds: 自定义哈希种子切片，长度必须等于depth
//
// 返回：
//
//	*CountMinSketch[T]: 初始化后的CM Sketch实例
//
// 注意：
//   - 若种子长度与depth不匹配，会触发panic（确保每个哈希函数有唯一种子）
//   - 内部会复制种子切片，避免外部修改影响内部哈希逻辑（防御性编程）
func NewCountMinSketchWithSeeds[T comparable](width, depth uint64, seeds []uint32) *CountMinSketch[T] {
	// 基础参数校验
	if width == 0 || depth == 0 {
		panic("width and depth must be positive")
	}
	// 种子长度校验：必须与行数（哈希函数个数）一致，否则哈希逻辑不完整
	if uint64(len(seeds)) != depth {
		panic(fmt.Sprintf("seeds length (%d) must equal depth (%d)", len(seeds), depth))
	}

	// 复制种子切片：避免外部修改seeds切片时影响内部哈希函数（值传递，隔离外部依赖）
	internalSeeds := make([]uint32, depth)
	copy(internalSeeds, seeds)

	// 初始化原子计数器矩阵（同NewCountMinSketch）
	counters := make([][]atomic.Uint64, depth)
	for i := range counters {
		counters[i] = make([]atomic.Uint64, width)
	}

	// 组装并返回实例
	return &CountMinSketch[T]{
		width:    width,
		depth:    depth,
		seeds:    internalSeeds,
		counters: counters,
	}
}

// Move 实现“移动构造”语义：将另一个CM Sketch的资源转移到当前实例
// 参数：
//
//	other: 待转移资源的CM Sketch实例（转移后other不可再使用）
//
// 作用：
//   - 避免资源拷贝（如大矩阵的内存复制），提升性能
//   - 常用于临时实例的资源复用（如分布式节点接收的Sketch合并后，释放临时实例资源）
//
// 注意：
//   - 若当前实例与other是同一对象，直接返回（避免自转移）
//   - 转移后会重置other的所有字段（避免资源重复释放或非法访问）
func (cms *CountMinSketch[T]) Move(other *CountMinSketch[T]) {
	// 避免自转移（无意义且可能导致资源混乱）
	if cms == other {
		return
	}
	// 转移核心资源：列数、行数、种子、计数器矩阵
	cms.width = other.width
	cms.depth = other.depth
	cms.seeds = other.seeds
	cms.counters = other.counters
	// 重置原实例：防止other后续被误用（如重复插入导致计数器混乱）
	other.width = 0
	other.depth = 0
	other.seeds = nil
	other.counters = nil
}

// hashToCol 计算指定元素在某一行（哈希函数）对应的列索引
// 核心逻辑：通过“种子+元素”的哈希值，将元素映射到[0, width-1]的列范围
// 参数：
//
//	item: 待映射的元素（任意可比较类型）
//	rowIdx: 目标行索引（对应第rowIdx个哈希函数）
//
// 返回：
//
//	uint64: 元素在当前行的列索引
//
// 注意：
//   - 会先做防御性检查，避免rowIdx越界（防止程序崩溃）
//   - 元素通过fmt.Sprintf转为字符串，复杂类型可自定义序列化方式优化性能
func (cms *CountMinSketch[T]) hashToCol(item T, rowIdx uint64) uint64 {
	// 防御性检查：确保行索引在有效范围内（避免访问seeds或counters时越界）
	if rowIdx >= cms.depth || rowIdx >= uint64(len(cms.seeds)) {
		panic(fmt.Sprintf("row index %d out of range (depth=%d, seeds length=%d)",
			rowIdx, cms.depth, len(cms.seeds)))
	}

	// 创建FNV-1a哈希函数实例：轻量级非加密哈希，速度快、碰撞率低，适合频率统计场景
	hasher := fnv.New64a()

	// 步骤1：将当前行的种子写入哈希函数（确保每行哈希逻辑独立）
	// 种子是32位，需转为4字节数组（BigEndian保证字节序一致性，跨平台兼容）
	seedBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(seedBytes, cms.seeds[rowIdx])
	_, _ = hasher.Write(seedBytes) // 忽略错误：seedBytes长度固定，不会写入失败

	// 步骤2：将元素转为字节并写入哈希函数（哈希函数仅能处理字节流）
	// 注意：fmt.Sprintf("%v", item)是通用方案，复杂类型（如自定义结构体）建议实现MarshalBinary优化
	itemBytes := []byte(fmt.Sprintf("%v", item))
	_, _ = hasher.Write(itemBytes) // 忽略错误：itemBytes是合法字符串转来的字节数组

	// 步骤3：计算哈希值并取模，将结果映射到[0, width-1]的列范围
	return hasher.Sum64() % cms.width
}

// Insert 插入元素到CM Sketch（线程安全，无锁设计）
// 参数：
//
//	item: 待插入的元素（任意可比较类型）
//
// 核心逻辑：
//  1. 对每个哈希函数（行），计算元素对应的列索引
//  2. 原子递增对应列的计数器（确保多线程并发安全）
//
// 性能：
//   - 时间复杂度O(d)（d为行数/哈希函数个数），通常d取4~10，接近O(1)
//   - 内存序使用Relaxed：计数器操作是独立的，无需与其他内存操作同步，最大化性能
func (cms *CountMinSketch[T]) Insert(item T) {
	// 遍历所有行（每个行对应一个哈希函数）
	for rowIdx := uint64(0); rowIdx < cms.depth; rowIdx++ {
		// 计算当前行的列索引
		colIdx := cms.hashToCol(item, rowIdx)
		// 原子递增计数器：Relaxed内存序足够（无需强同步，计数器操作独立）
		cms.counters[rowIdx][colIdx].Add(1)
	}
}

// Count 估算元素的出现频率（返回计数器最小值，线程安全）
// 参数：
//
//	item: 待估算频率的元素（任意可比较类型）
//
// 返回：
//
//	uint64: 元素的估算频率（是真实频率的上界，即估算值≥真实值）
//
// 核心逻辑：
//  1. 对每个哈希函数（行），计算元素对应的列索引
//  2. 读取所有对应列的计数器值
//  3. 返回最小值（哈希碰撞会导致计数器高估，最小值最接近真实值）
func (cms *CountMinSketch[T]) Count(item T) uint64 {
	// 初始化最小值为uint64的最大值：确保第一个计数器值能覆盖它
	minCount := uint64(math.MaxUint64)
	// 遍历所有行，读取对应列的计数器值
	for rowIdx := uint64(0); rowIdx < cms.depth; rowIdx++ {
		colIdx := cms.hashToCol(item, rowIdx)
		// 原子加载计数器值：Relaxed内存序足够（仅读取，无需同步其他操作）
		count := cms.counters[rowIdx][colIdx].Load()
		// 更新最小值（找最接近真实值的上界）
		if count < minCount {
			minCount = count
		}
	}
	// 返回估算频率（若元素未插入，返回0，因所有计数器初始为0）
	return minCount
}

// Clear 重置所有计数器为0（线程安全，不影响并发插入）
// 核心逻辑：
//
//	遍历所有行和列，用原子操作将计数器设为0
//
// 注意：
//   - 即使有并发Insert操作，原子Store也能保证计数器最终被设为0（可能存在短暂中间状态，但最终正确）
//   - 适用于周期性统计场景（如每天凌晨重置统计数据）
func (cms *CountMinSketch[T]) Clear() {
	// 遍历所有行和列
	for rowIdx := uint64(0); rowIdx < cms.depth; rowIdx++ {
		for colIdx := uint64(0); colIdx < cms.width; colIdx++ {
			// 原子存储0：覆盖原有值，确保线程安全
			cms.counters[rowIdx][colIdx].Store(0)
		}
	}
}

// Merge 合并两个兼容的CM Sketch，返回新实例（原实例不变，线程安全）
// 兼容条件：
//  1. 列数（width）相同
//  2. 行数（depth）相同
//  3. 哈希种子（seeds）完全相同
//
// 参数：
//
//	other: 待合并的另一个CM Sketch实例
//
// 返回：
//
//	*CountMinSketch[T]: 合并后的新实例
//	error: 若不兼容，返回错误信息（如width/depth不匹配、种子不同）
//
// 核心逻辑：
//  1. 兼容性检查：确保两个Sketch的哈希逻辑一致
//  2. 创建新实例：复用原实例的width、depth、seeds（保证合并后逻辑不变）
//  3. 计数器累加：对应位置的计数器值相加（等价于两个流拼接后插入单个Sketch）
func (cms *CountMinSketch[T]) Merge(other *CountMinSketch[T]) (*CountMinSketch[T], error) {
	// 兼容性检查1：列数和行数必须相同（否则矩阵结构不匹配，无法对应累加）
	if cms.width != other.width || cms.depth != other.depth {
		return nil, fmt.Errorf("incompatible width/depth: cms(%d×%d) vs other(%d×%d)",
			cms.width, cms.depth, other.width, other.depth)
	}
	// 兼容性检查2：种子长度必须相同（避免哈希函数个数不匹配）
	if len(cms.seeds) != len(other.seeds) {
		return nil, fmt.Errorf("seeds length mismatch: cms(%d) vs other(%d)",
			len(cms.seeds), len(other.seeds))
	}
	// 兼容性检查3：每个种子必须完全相同（确保哈希逻辑一致，元素映射位置相同）
	for i := range cms.seeds {
		if cms.seeds[i] != other.seeds[i] {
			return nil, fmt.Errorf("incompatible seed at index %d", i)
		}
	}

	// 创建合并后的新实例：使用指定种子构造，确保与原实例逻辑一致
	merged := NewCountMinSketchWithSeeds[T](cms.width, cms.depth, cms.seeds)

	// 计数器累加：遍历所有(row, col)，将两个实例的计数器值相加后存入新实例
	for rowIdx := uint64(0); rowIdx < cms.depth; rowIdx++ {
		for colIdx := uint64(0); colIdx < cms.width; colIdx++ {
			// 原子加载两个实例的计数器值（无并发风险，因仅读取）
			cmsCount := cms.counters[rowIdx][colIdx].Load()
			otherCount := other.counters[rowIdx][colIdx].Load()
			// 原子存储累加结果到新实例
			merged.counters[rowIdx][colIdx].Store(cmsCount + otherCount)
		}
	}

	// 返回合并后的新实例（原实例未被修改，可继续使用）
	return merged, nil
}

// TopK 从候选集中返回估算频率最高的k个元素（线程安全）
// 注意：
//
//	CM Sketch不存储元素本身，需传入候选集（如已知的所有可能元素）
//
// 参数：
//
//	k: 需返回的Top元素个数
//	candidates: 元素候选集（所有可能需要统计的元素）
//
// 返回：
//
//	[]T: 估算频率从高到低的前k个元素（若候选集不足k个，返回全部）
func (cms *CountMinSketch[T]) TopK(k uint64, candidates []T) []T {
	// 步骤1：定义结构体存储“元素-频率”对（用于后续排序）
	type itemCount struct {
		item  T      // 元素本身
		count uint64 // 元素的估算频率
	}
	// 初始化切片：容量为候选集长度，避免动态扩容开销
	icList := make([]itemCount, 0, len(candidates))
	// 遍历候选集，计算每个元素的估算频率
	for _, item := range candidates {
		icList = append(icList, itemCount{
			item:  item,
			count: cms.Count(item), // 调用Count方法获取估算频率
		})
	}

	// 步骤2：排序：先按频率降序，频率相同则按元素字符串序（保证排序稳定性）
	sort.Slice(icList, func(i, j int) bool {
		// 首要排序条件：频率高的在前
		if icList[i].count != icList[j].count {
			return icList[i].count > icList[j].count
		}
		// 次要排序条件：频率相同时，按元素字符串的字典序（避免排序结果随机）
		return fmt.Sprintf("%v", icList[i].item) < fmt.Sprintf("%v", icList[j].item)
	})

	// 步骤3：取前k个元素（处理候选集不足k个的情况）
	take := uint64(len(icList)) // 初始取全部候选集
	if k < take {               // 若k小于候选集长度，取前k个
		take = k
	}
	// 初始化结果切片：容量为take，避免动态扩容
	topK := make([]T, 0, take)
	for i := uint64(0); i < take; i++ {
		topK = append(topK, icList[i].item)
	}

	// 返回TopK结果
	return topK
}
